package com.oliver.oliveraiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.oliver.oliveraiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @ClassName ToolsCallAgent
 * @Description TODO
 * @Author snow
 * @Date 2025/5/24 0:49
 **/

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{

    //  可用的工具
    private ToolCallback[] availableTools;

    //  工具调用管理者
    private ToolCallingManager toolCallingManager;

    //  工具调用结果
    private ChatResponse toolCallChatResponse;

    //  禁用内置的工具调用机制，自己维护上下文
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools){
        super();
        this.availableTools = availableTools;
        toolCallingManager = ToolCallingManager.builder().build();
        //  禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 处理当前状态并决定下一步行动
     * @return
     */
    @Override
    public boolean think() {
        if (StrUtil.isNotBlank(getNextStepPrompt())){
            getMessageList().add(new UserMessage(getNextStepPrompt()));
        }

        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(messageList, chatOptions);
        try{
            //  获取带工具选项的响应
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            //  记录响应，用于 Act 操作
            this.toolCallChatResponse = chatResponse;
            //  获取助手消息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //  输出提示消息
            String result = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            log.info(getName() + "的思考" + result);
            log.info(getName() + "选择了" + toolCallList.size() + "个工具来使用");
            //  工具调用详情
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s,参数:%s",
                            toolCall.name(), toolCall.arguments())
                    ).collect(Collectors.joining("\n"));
            log.info("工具调用的详情信息：" +toolCallInfo);

            //  只有不工具调用时，才记录助手消息
            if (toolCallList.isEmpty()){
                getMessageList().add(assistantMessage);
                return false;
            }else {
                //  需要调用工具，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        }catch (Exception e){
            e.printStackTrace();
            log.error(getName() + "的思考出错：" + e.getMessage());
            getMessageList().add(new AssistantMessage("思考出错：" + e.getMessage()));
            return false;
        }

    }

    /**
     * 执行工具调用并返回处理结果
     * @return
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()){
            return "无需调用工具";
        }

        //  调用工具
        Prompt prompt = new Prompt(getMessageList(),chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        //  记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        setMessageList(toolExecutionResult.conversationHistory());

        //  当前工具调用的结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具" + response.name() + "完成了它的任务！结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        log.info("工具调用结果 : " + results);

        //  判断是否有执行终止工具，如果有，则修改 agent 的状态
        boolean terminateCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));

        if (terminateCalled) {
            setState(AgentState.FINISHED);
        }

        return results;
    }
}
